{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3, Topic 1, Lab B: AES256 Bootloader Attack (MAIN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "NOTE: This lab references some (commercial) training material on [ChipWhisperer.io](https://www.ChipWhisperer.io). You can freely execute and use the lab per the open-source license (including using it in your own courses if you distribute similarly), but you must maintain notice about this source location. Consider joining our training course to enjoy the full experience.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *Through the previous labs, we've gained a lot of tools to attack unknown embedded devices: SPA, DPA, CPA, trace resynchronization, and more. In this lab, we'll be using some of those techniques to break a more realistic target: a bootloader. Note that there are two versions of this lab. In this one (Lab B), we'll start with no information that couldn't be revealed by watching code be sent to the bootloader. Everything else, we'll need to figure out for ourselves, such as what encryption algorithm the target is using, how it's using it, etc. In lab A, this information will be given and you'll just focus on the attack. It's up to you whether you want to run this lab or Lab A!*\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "\n",
    "* Observing power traces to figure out what encryption operation it's running\n",
    "* Applying CPA and DPA to break different parts of the bootloader\n",
    "* Understanding different operating modes for block ciphers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "In the world of microcontrollers, a bootloader is a special piece of firmware that is made to let the user upload new programs into memory. This is especially useful for devices with complex code that may need to be patched or otherwise updated in the future - a bootloader makes it possible for the user to upload a patched version of the firmware onto the micro. The bootloader receives information from a communication line (a USB port, serial port, ethernet port, WiFi connection, etc...) and stores this data into program memory. Once the full firmware has been received, the micro can happily run its updated code.\n",
    "\n",
    "There is one big security issue to worry about with bootloaders. A company may want to stop their customers from writing their own firmware and uploading it onto the micro. For example, this might be for protection reasons - hackers might be able to access parts of the device that weren't meant to be accessed. One way of stopping this is to add encryption. The company can add their own secret signature to the firmware code and encrypt it with a secret key. Then, the bootloader can decrypt the incoming firmware and confirm that the incoming firmware is correctly signed. Users will not know the secret key or the signature tied to the firmware, so they won't be able to \"fake\" their own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \n",
    "cd ../../../firmware/mcu/bootloader-aes256\n",
    "make PLATFORM=$1 CRYPTO_TARGET=NONE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../../Setup_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fw_path = \"../../../firmware/mcu/bootloader-aes256/bootloader-aes256-{}.hex\".format(PLATFORM)\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Situation\n",
    "\n",
    "Simply put, we've got a target device running an encrypted bootloader (a program used to upload new code onto a device) and we want to see if we can get our own code running on the device. We've done a bit of sniffing on the serial lines when the device's firmware is being updated and we've learned the following:\n",
    "\n",
    "* The device communicates over serial at 38400bps\n",
    "* When writing memory, the first byte is always zero (probably a command byte)\n",
    "* There's a 16 byte block of random looking memory (aka it doesn't look like firmware). This part is probably encrypted\n",
    "* There's a 2 byte CRC at the end of each message\n",
    "* There's no repetition in the ciphertext.\n",
    "\n",
    "All together this looks like:\n",
    "\n",
    "```\n",
    "       |<-------- Encrypted block (16 bytes) ---------->|\n",
    "       |                                                |\n",
    "+------+------+------+------+------+------+ .... +------+------+------+\n",
    "| 0x00 |              Random looking data               |   CRC-16    |\n",
    "+------+------+------+------+------+------+ .... +------+------+------+\n",
    "```\n",
    "\n",
    "After sending data to the bootloader it responds with either `0xA4` or `0xA1`. The former only happened when we sent a bad CRC.\n",
    "\n",
    "This time, we won't be triggering off of our trigger pins (you can remove them from the code if you'd like).\n",
    "\n",
    "From our initial sniffing of the communication lines, we've got the first few messages that were sent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(\"./firmware.pickle\", \"rb\") as f:\n",
    "    encrypted_firmware = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Doing Recon\n",
    "\n",
    "Our first step will be to see if we can learning anything about the bootloader from looking at its power traces. Let's start with the boot sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.triggers = \"nrst\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.adc.samples = 24400"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.arm()\n",
    "reset_target(scope)\n",
    "scope.capture()\n",
    "wave = scope.get_last_trace()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(wave)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The device does appear to be doing something, but it's clearly nothing major - no encryptions or anything. Every microcontroller has boot code and operations that run when it's reset. The device may even have its own bootloader running in ROM! Let's move onto the messages themselves. We do know that there's a CRC for data integrety. We can use the following code to calculate the CRC for us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Class Crc\n",
    "#############################################################\n",
    "# These CRC routines are copy-pasted from pycrc, which are:\n",
    "# Copyright (c) 2006-2013 Thomas Pircher <tehpeh@gmx.net>\n",
    "#\n",
    "class Crc(object):\n",
    "    \"\"\"\n",
    "    A base class for CRC routines.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, width, poly):\n",
    "        \"\"\"The Crc constructor.\n",
    "\n",
    "        The parameters are as follows:\n",
    "            width\n",
    "            poly\n",
    "            reflect_in\n",
    "            xor_in\n",
    "            reflect_out\n",
    "            xor_out\n",
    "        \"\"\"\n",
    "        self.Width = width\n",
    "        self.Poly = poly\n",
    "\n",
    "\n",
    "        self.MSB_Mask = 0x1 << (self.Width - 1)\n",
    "        self.Mask = ((self.MSB_Mask - 1) << 1) | 1\n",
    "\n",
    "        self.XorIn = 0x0000\n",
    "        self.XorOut = 0x0000\n",
    "\n",
    "        self.DirectInit = self.XorIn\n",
    "        self.NonDirectInit = self.__get_nondirect_init(self.XorIn)\n",
    "        if self.Width < 8:\n",
    "            self.CrcShift = 8 - self.Width\n",
    "        else:\n",
    "            self.CrcShift = 0\n",
    "\n",
    "    def __get_nondirect_init(self, init):\n",
    "        \"\"\"\n",
    "        return the non-direct init if the direct algorithm has been selected.\n",
    "        \"\"\"\n",
    "        crc = init\n",
    "        for i in range(self.Width):\n",
    "            bit = crc & 0x01\n",
    "            if bit:\n",
    "                crc ^= self.Poly\n",
    "            crc >>= 1\n",
    "            if bit:\n",
    "                crc |= self.MSB_Mask\n",
    "        return crc & self.Mask\n",
    "\n",
    "\n",
    "    def bit_by_bit(self, in_data):\n",
    "        \"\"\"\n",
    "        Classic simple and slow CRC implementation.  This function iterates bit\n",
    "        by bit over the augmented input message and returns the calculated CRC\n",
    "        value at the end.\n",
    "        \"\"\"\n",
    "        # If the input data is a string, convert to bytes.\n",
    "        if isinstance(in_data, str):\n",
    "            in_data = [ord(c) for c in in_data]\n",
    "\n",
    "        register = self.NonDirectInit\n",
    "        for octet in in_data:\n",
    "            for i in range(8):\n",
    "                topbit = register & self.MSB_Mask\n",
    "                register = ((register << 1) & self.Mask) | ((octet >> (7 - i)) & 0x01)\n",
    "                if topbit:\n",
    "                    register ^= self.Poly\n",
    "\n",
    "        for i in range(self.Width):\n",
    "            topbit = register & self.MSB_Mask\n",
    "            register = ((register << 1) & self.Mask)\n",
    "            if topbit:\n",
    "                register ^= self.Poly\n",
    "\n",
    "        return register ^ self.XorOut\n",
    "    \n",
    "bl_crc = Crc(width = 16, poly=0x1021)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's definte a function to do the communication and capture a trace for us. We can try triggering off of our message. There's not much memory on the target, so it's probably decryption on the fly instead of reading in a whole bunch of memory, then doing the decryption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.triggers = \"tio2\"\n",
    "scope.adc.samples = 24400\n",
    "scope.adc.decimate = ??? # try to get the full encryption in a single trace, then set back to 1\n",
    "scope.adc.offset = 0\n",
    "def cap_trace(enc_block):\n",
    "    message = [0x00]\n",
    "    target.read()\n",
    "\n",
    "    key, text = ktp.next()\n",
    "    message.extend(enc_block)\n",
    "\n",
    "    crc = bl_crc.bit_by_bit(enc_block)\n",
    "    message.append(crc >> 8)\n",
    "    message.append(crc & 0xFF)\n",
    "\n",
    "    \n",
    "\n",
    "    target.write(message[:-1])\n",
    "    time.sleep(0.01)\n",
    "    scope.arm()\n",
    "    target.write([crc&0xFF])\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        print('Timeout happened during acquisition')\n",
    "    response = target.read()\n",
    "    if response:\n",
    "        if ord(response[0]) != 0xA4:\n",
    "            # Bad response, just skip\n",
    "            #print(\"Bad response: {:02X}\".format(ord(response[0])))\n",
    "            return None\n",
    "\n",
    "    return scope.get_last_trace()\n",
    "    \n",
    "ktp = cw.ktp.Basic()\n",
    "text, key = ktp.next()\n",
    "wave = cap_trace(text)\n",
    "wave2 = cap_trace(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(wave) * cw.plot(wave2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've found the decryption. Some things to notice from this:\n",
    "\n",
    "1. The target immediately goes from reading the ciphertext to running a decryption - there's no preprocessing done at all. This means the ciphertext is likely being immediately fed into the decryption algorithm.\n",
    "1. We can see operations in the encryption being repeated. This looks a lot like AES - we can see a long distinct operation that's probably MixColumns in there and the overall structure looks a lot like what we saw when looking at TINYAES128C encryptions. Let's make a weak assumption that this is AES128 - we can adjust this as we learn more about the bootloader.\n",
    "1. We can't see the full encryption\n",
    "1. There's some jitter here. We'll probably have to resync the traces if we run a CPA attack\n",
    "\n",
    "Since we can't see the full encryption, decimate the ADC (we don't care too much about the fine details here) and take another look...\n",
    "\n",
    "\n",
    "## The full encryption\n",
    "\n",
    "You should see that instead of 9 repititions of MixColumns (or what we assume is MixColumns), there's actually 13. This rules out AES128, but AES256 actually has 14 rounds! We can still attack AES256 without much issue: we basically just have to run two CPA attacks, one for the first half of the key and another for the second half. Again, we're not 100% about this, but it's a good starting point. \n",
    "\n",
    "As we mentioned in the debriefing about the bootloader, the ciphertext never seems to repeat. It's pretty unlikely that this device is using straight AES256 (if it even is using AES256), since firmware usually has blocks that repeat. More likely is that AES is being used as a stream cipher: https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation. This could pose an immediate problem for our attack efforts: we need to know either what's going into or coming out of the AES block to perform a CPA attack. However, some of the modes listed on that page (if it's even using one on that page) feed an IV or a counter into that block instead of the plaintext or the ciphertext. We didn't see any encryption operations happening on startup (which the device could've done if it was using one of these IV/counter modes), so we'll probably be okay with a normal CPA attack. Let's try it and see if we can get anything out of it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import trange\n",
    "project = cw.create_project(\"projects/Tutorial_A5\", overwrite=True)\n",
    "scope.adc.offset = 0\n",
    "scope.adc.decimate = 1\n",
    "for i in trange(100):\n",
    "    ktp = cw.ktp.Basic()\n",
    "    key, text = ktp.next()\n",
    "    wave = cap_trace(text)\n",
    "    trace = cw.Trace(wave, text, bytearray([0]*16), bytearray([0]*16))\n",
    "    project.traces.append(trace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(project.waves[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can eliminate a lot of this jitter by using the resync SAD module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "\n",
    "leak_model = cwa.leakage_models.inverse_sbox_output\n",
    "resync = cwa.preprocessing.ResyncSAD(project)\n",
    "resync.enabled=True\n",
    "resync.ref_trace = 0\n",
    "resync.target_window = (???, ???)\n",
    "resync.max_shift = 7000\n",
    "new_proj = resync.preprocess()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check to make sure your traces are resynced:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt = cw.plot([])\n",
    "for i in range(10):\n",
    "    plt *= cw.plot(new_proj.waves[i])\n",
    "\n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All that's left is to actually run the attack. We don't know the correct key, so it won't be highlighted in red."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack = cwa.cpa(new_proj, leak_model)\n",
    "#attack.pont_range = [???, ???]\n",
    "\n",
    "#key = [0xea, 0x79, 0x79, 0x20, 0xc8, 0x71, 0x44, 0x7d, 0x46, 0x62, 0x5f, 0x51, 0x85, 0xc1, 0x3b, 0xcb]\n",
    "\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "attack_results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The difference in correlation between the best key guess and the next best one makes this look very promising! We now know we're correct about two things:\n",
    "\n",
    "1. The bootloader is actually decrypting the ciphertext\n",
    "1. The device is using AES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calc_round_key = attack_results.key_guess()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that done, we now need to get the second half of the key. Pop over to [Extending AES-128 Attacks to AES-256](Extending%20AES-128%20Attacks%20to%20AES-256.ipynb), since that page explains how to do that...\n",
    "\n",
    "Back? Let's go through and see if that theory actually holds up.\n",
    "\n",
    "To make a new model, we start off by inheriting the `AESLeakageHelper` class. We need to make a `leakage()` method that calculates the Hamming weight we use in the CPA attack. To get you started:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "class AES256_Round13_Model(cwa.AESLeakageHelper):\n",
    "    def leakage(self, pt, ct, guess, bnum):\n",
    "        #You must put YOUR recovered 14th round key here - this example may not be accurate!\n",
    "        calc_round_key = [0xea, 0x79, 0x79, 0x20, 0xc8, 0x71, 0x44, 0x7d, 0x46, 0x62, 0x5f, 0x51, 0x85, 0xc1, 0x3b, 0xcb]\n",
    "        state = reverse_round_14(self, pt, calc_round_key)\n",
    "        state = reverse_round_13(self, state) #reverse state just before inv_subbytes\n",
    "        return self.inv_sbox(state[bnum] ^ guess[bnum])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we just need to make the `reverse_round_14()` and `reverse_round_13()` functions. By passing the class in, we get access to `self.inv_shiftrows()`, `self.inv_subbytes()`, and `self.inv_mixcolumns()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reverse_round_14(self, pt, key):\n",
    "    state = [pt[i] ^ key[i] for i in range(16)] #AddRoundKey\n",
    "    state = ???\n",
    "    state = ???\n",
    "    return state # we're now at the end of decryption round 1\n",
    "\n",
    "def reverse_round_13(self, state):\n",
    "    state = ???\n",
    "    state = ???\n",
    "    return state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From there, just update the leakage model and rerun the attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model = cwa.leakage_models.new_model(AES256_Round13_Model)\n",
    "attack.leak_model = leak_model\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "attack_results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you built your model correctly, you should again see a pretty likely guess for a key. All that's left now is to combine the 14th and 13th round keys, then use that to figure out the 0th and 1st round keys.\n",
    "\n",
    "We'll start by getting the transformed 13th round key out of the attack results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calc_round_key = [0xea, 0x79, 0x79, 0x20, 0xc8, 0x71, 0x44, 0x7d, 0x46, 0x62, 0x5f, 0x51, 0x85, 0xc1, 0x3b, 0xcb]\n",
    "rec_key = calc_round_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_key2 = []\n",
    "for bnum in attack_results.find_maximums():\n",
    "    print(\"Best Guess = 0x{:02X}, Corr = {}\".format(bnum[0][0], bnum[0][2]))\n",
    "    rec_key2.append(bnum[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to transform that key into the actual 13th round key by running it through ShiftRows and MixColumns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_key2 = cwa.aes_funcs.shiftrows(rec_key2)\n",
    "real_key2 = cwa.aes_funcs.mixcolumns(real_key2)\n",
    "\n",
    "print(\"Recovered:\", end=\"\")\n",
    "for subkey in real_key2:\n",
    "    print(\" {:02X}\".format(subkey), end=\"\")\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "then we can combine the keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_key_comb = real_key2.copy()\n",
    "rec_key_comb.extend(rec_key)\n",
    "\n",
    "print(\"Key:\", end=\"\")\n",
    "for subkey in rec_key_comb:\n",
    "    print(\" {:02X}\".format(subkey), end=\"\")\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and use ChipWhisperer's built in key scheduler to reverse them to the 0th and 1st round keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btldr_key = leak_model.key_schedule_rounds(rec_key_comb, 13, 0)\n",
    "btldr_key.extend(leak_model.key_schedule_rounds(rec_key_comb, 13, 1))\n",
    "print(\"Key:\", end=\"\")\n",
    "for subkey in btldr_key:\n",
    "    print(\" {:02X}\".format(subkey), end=\"\")\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we were clearly right about the bootloader running AES256! However, if we try decrypting some of our firmware with the encryption key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Crypto.Cipher import AES\n",
    "cipher = AES.new(bytes(btldr_key), AES.MODE_ECB)\n",
    "print(bytearray(cipher.decrypt(encrypted_firmware[:16])))\n",
    "print(bytearray(cipher.decrypt(encrypted_firmware[16:32])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we expected, we still get gibberish out of this - the target is definitely using AES as a stream cipher. The question now is, which block cipher mode is it using? Well, we know the ciphertext is being decrypted. We can narrow it down a bit by looking at the end of the encryption. Increase the offset until you reach the end of the encryption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#scope.adc.offset = 50000\n",
    "#wave = cap_trace(text)\n",
    "#key, text = ktp.next()\n",
    "wave = cap_trace(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.plot(wave)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It might be hard to pick out, but you should be able to find 16 XOR short operations just after the end of the last round of AES. It might be using CBC mode, which means it'll be using the ciphertext as part of the encryption and decryption of subsequent blocks. Let's do a quick check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cipher = AES.new(bytes(btldr_key), AES.MODE_ECB)\n",
    "dec_fw = cipher.decrypt(encrypted_firmware[16:32])\n",
    "fw = [dec_fw[i] ^ encrypted_firmware[i] for i in range(16)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fw)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like we're finally getting some valid output. Let's try the next block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cipher = AES.new(bytes(btldr_key), AES.MODE_ECB)\n",
    "dec_fw = cipher.decrypt(encrypted_firmware[32:48])\n",
    "fw = [dec_fw[i] ^ encrypted_firmware[i+16] for i in range(16)]\n",
    "print(fw)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is very promising! This firmware is repeated twice. The all `FFs` is probably just empty flash memory. The beginning is more curious though:\n",
    "\n",
    "* It's probably not a flash constant - otherwise it wouldn't be in two blocks in a row\n",
    "* That's only enough room for a few instructions at most. Again, it's a little strange that it would be repeated like this\n",
    "\n",
    "It might be some sort of signature. After all, the device doesn't want to write anything to memory unless the ciphertext has properly been decrypted.\n",
    "\n",
    "There's still the issue of the first block of memory though. There's another secret value called an initialization vector that we need to decrypt that block. To be able to recover that, we'll need to revisit the DPA attack:\n",
    "\n",
    "### Attack Theory\n",
    "\n",
    "The bootloader applies the IV to the AES decryption result (`DR`) by calculating\n",
    "\n",
    "\n",
    "$\\text{PT} = \\text{DR} \\oplus \\text{IV}$\n",
    "\n",
    "where DR is the decrypted ciphertext, IV is the secret vector, and PT is the plaintext that the bootloader will use later. We only have access to one of these: since we know the AES-256 key, we can calculate DR. This exclusive or will be visible in the power traces.\n",
    "\n",
    "This is enough information for us to attack a single bit of the IV. Suppose we only wanted to get the first bit (bit 0) of the first byte (byte 0) of the IV. We could do the following:\n",
    "\n",
    "* Split all of the traces into two groups: those with `(DR[0] & 0x01) = 0`, and those with `(DR[0] & 0x01) = 1`. \n",
    "* Calculate the average trace for both groups.\n",
    "* Find the difference between the two averages. Provided we've got sufficient data, we should see a spike where the xor is occuring.\n",
    "* Look at the direction of the spike to decide if the IV bit is 0 `(PT[0] = DR[0])` or if the IV bit is 1 `(PT[0] = ~DR[0])`.\n",
    "\n",
    "This is effectively a DPA attack on a single bit of the IV. We can repeat this attack across the whole IV by instead separating by `(DR[byte] & (1 << bit) = 0` and `(DR[byte] & (1 << bit) = bit`.\n",
    "\n",
    "We'll need to reset the device every encryption since it only uses the IV in the first encryption. This leads to a slightly modified capture loop. You'll need to adjust your offset since we're now triggering near the beginning of the UART transmit instead of nera the end. Run the loop, then interrupt it to get a wave. Then plot and adjust your offset. Repeat until you're near the end of the encryption again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Crypto.Cipher import AES\n",
    "import numpy as np\n",
    "\n",
    "from tqdm.notebook import trange\n",
    "import numpy as np\n",
    "import time\n",
    "traces = []\n",
    "keys = []\n",
    "plaintexts = []\n",
    "\n",
    "from tqdm.notebook import trange\n",
    "project = cw.create_project(\"projects/Tutorial_A5_IV\", overwrite=True)\n",
    "scope.adc.decimate = 1\n",
    "scope.adc.offset = 51000\n",
    "scope.adc.timeout = 1\n",
    "scope.trigger.triggers = \"tio2\"\n",
    "for i in trange(1000):\n",
    "    scope.io.nrst = 0\n",
    "    time.sleep(0.02)\n",
    "    scope.io.nrst = \"high_z\"\n",
    "    time.sleep(0.01)\n",
    "    okay = 0\n",
    "    while not okay:\n",
    "        target.write('\\0xxxxxxxxxxxxxxxxxx')\n",
    "        time.sleep(0.005)\n",
    "        response = target.read()\n",
    "        i += 1\n",
    "        if response:\n",
    "            if ord(response[0]) == 0xA1:\n",
    "                okay = 1\n",
    "    message = [0x00]\n",
    "    \n",
    "    target.flush()\n",
    "    \n",
    "    key, text = ktp.new_pair()  # manual creation of a key, text pair can be substituted here\n",
    "    \n",
    "    wave = cap_trace(text)\n",
    "    if wave is None:\n",
    "        continue\n",
    "    \n",
    "    \n",
    "    #wave = scope.get_last_trace()\n",
    "    trace = cw.Trace(wave, text, bytearray([0]*16), bytearray([0]*16))\n",
    "    project.traces.append(trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, we'll again need to resync to get rid of the jitter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt = cw.plot([])\n",
    "for i in range(2):\n",
    "    plt *= cw.plot(project.waves[i])\n",
    "    \n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adjust the target window here as necessary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "\n",
    "leak_model = cwa.leakage_models.inverse_sbox_output\n",
    "resync = cwa.preprocessing.ResyncSAD(project)\n",
    "resync.enabled=True\n",
    "resync.ref_trace = 0\n",
    "resync.target_window = (???, ???)\n",
    "resync.max_shift = 6000\n",
    "new_proj = resync.preprocess()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt = cw.plot([])\n",
    "for i in range(10):\n",
    "    plt *= cw.plot(new_proj.waves[i])\n",
    "    \n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some numpy functions will be useful here, so we'll convert our ChipWhisperer project to numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_array = np.array([new_proj.waves[i] for i in range(len(new_proj.traces))])\n",
    "textin_array = np.array([new_proj.textins[i] for i in range(len(new_proj.traces))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to decrypt what we sent to the device to get one half of the input to the XOR:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "knownkey = [0x94, 0x28, 0x5D, 0x4D, 0x6D, 0xCF, 0xEC, 0x08, 0xD8, 0xAC, 0xDD, 0xF6, 0xBE, 0x25, 0xA4, 0x99,\n",
    "            0xC4, 0xD9, 0xD0, 0x1E, 0xC3, 0x40, 0x7E, 0xD7, 0xD5, 0x28, 0xD4, 0x09, 0xE9, 0xF0, 0x88, 0xA1]\n",
    "\n",
    "knownkey = bytes(knownkey)\n",
    "dr = []\n",
    "aes = AES.new(knownkey, AES.MODE_ECB)\n",
    "for i in range(len(new_proj.traces)):\n",
    "    ct = bytes(textin_array[i])\n",
    "    pt = aes.decrypt(ct)\n",
    "    d = [bytearray(pt)[i] for i in range(16)]\n",
    "    dr.append(d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The basic idea is the same as the DPA attack: guess a bit and group traces based on that. We'll do the first byte here as an example. For each of the bits, you should see roughly half the traces fall into each group:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-1-b8a799d34f4a>, line 2)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-1-b8a799d34f4a>\"\u001b[1;36m, line \u001b[1;32m2\u001b[0m\n\u001b[1;33m    byte =\u001b[0m\n\u001b[1;37m           ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "grouped_byte_traces = []\n",
    "byte = ??? # start with zero, then come back later and change\n",
    "\n",
    "for bit in range(8):\n",
    "    grouped_bit_traces = [], []\n",
    "    for i in range(len(new_proj.traces)):\n",
    "        if (dr[i][byte] & (1 << bit)):\n",
    "            grouped_bit_traces[0].append(trace_array[i])\n",
    "        else:\n",
    "            grouped_bit_traces[1].append(trace_array[i])\n",
    "    grouped_byte_traces.append(grouped_bit_traces)\n",
    "    print(len(grouped_bit_traces[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to figure out where the XOR operation is happening. You'll need to use the shape of the plot and the following plot of the difference of means for each bit. Each colour represents a different bit. For the correct part on the plot, you should see a distinct separation, with some bits peaking above zero, and others peaking below zero. If you see some colours in between peaks, it is probably not the right location. Repeat this for a few bytes - the location should change. Note down this change, as you'll have to use it to adjust the analysis location later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find averages and differences\n",
    "diffs = []\n",
    "for i in range(8):\n",
    "    means = np.average(grouped_byte_traces[i][0], axis=0), np.average(grouped_byte_traces[i][1], axis=0)\n",
    "    diffs.append(means[1] - means[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split traces into 2 groups\n",
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "\n",
    "output_notebook()\n",
    "p = figure()\n",
    "\n",
    "xrange = range(len(diffs[0]))\n",
    "xrange2 = range(len(trace_array[0]))\n",
    "colours = [\"red\", \"blue\", \"green\", \"black\"]\n",
    "plt = cw.plot([])\n",
    "for i in range(8):\n",
    "    plt *= cw.plot(diffs[i]).opts(color=colours[i%4])\n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fill in the position of the XOR, as well as how much it changes for each byte. All we're doing here is going byte by byte and bit by bit, and seeing if the difference in means is greater than or less than zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btldr_IV = [0] * 16 #181\n",
    "for byte in range(16):\n",
    "    location = ??? + byte * ???\n",
    "    iv = 0\n",
    "    for bit in range(8):\n",
    "        pt_bits = [((dr[i][byte] >> (7-bit)) & 0x01) for i in range(len(new_proj.traces))]\n",
    "\n",
    "        # Split traces into 2 groups\n",
    "        groupedPoints = [[] for _ in range(2)]\n",
    "        for i in range(len(new_proj.traces)):\n",
    "            groupedPoints[pt_bits[i]].append(trace_array[i][location])\n",
    "            \n",
    "        means = []\n",
    "        for i in range(2):\n",
    "            means.append(np.average(groupedPoints[i]))\n",
    "        diff = means[1] - means[0]\n",
    "        \n",
    "        iv_bit = 1 if diff > 0 else 0\n",
    "        iv = (iv << 1) | iv_bit\n",
    "        \n",
    "        print(iv_bit, end = \" \")\n",
    "        \n",
    "    print(\"{:02X}\".format(iv))\n",
    "    btldr_IV[byte] = iv\n",
    "    \n",
    "print(btldr_IV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can do the full decryption!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cipher = AES.new(bytes(btldr_key), AES.MODE_ECB)\n",
    "first_pt = cipher.decrypt(encrypted_firmware[:16])\n",
    "first_pt = [first_pt[i] ^ btldr_IV[i] for i in range(16)]\n",
    "print(bytearray(first_pt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the first line of firmware is `deadbeefaabbccddeeff0011`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
